package city.spring.configure.security;

import org.springframework.security.jwt.JwtHelper;
import org.springframework.security.jwt.crypto.sign.MacSigner;
import org.springframework.security.jwt.crypto.sign.RsaSigner;
import org.springframework.security.jwt.crypto.sign.Signer;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.util.JsonParser;
import org.springframework.security.oauth2.common.util.JsonParserFactory;
import org.springframework.security.oauth2.common.util.RandomValueStringGenerator;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

import java.security.KeyPair;
import java.security.interfaces.RSAPrivateKey;
import java.util.Collections;
import java.util.Map;

/**
 * 自定义JWT Token转换器，加入自定义头部，为了适应JWK的kid参数，使资源服务能够使用JWK方式来验证 jwt Token 信息
 *
 * @author HouKunLin
 * @date 2019/12/11 0011 16:05
 */
public class CustomJwtAccessTokenConverter extends JwtAccessTokenConverter {
    private final JsonParser objectMapper = JsonParserFactory.create();
    private final Map<String, String> customHeaders;
    private String verifierKey = new RandomValueStringGenerator().generate();
    private Signer signer = new MacSigner(verifierKey);

    public CustomJwtAccessTokenConverter() {
        super();
        // 自定义 jwt 头部信息，添加kid字段信息，使其可以使用JWK方式校验Token
        this.customHeaders = Collections.singletonMap("kid", "oauth-jwk-key-id");
    }

    @Override
    protected String encode(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        String content;
        try {
            content = this.objectMapper.formatMap(getAccessTokenConverter().convertAccessToken(accessToken, authentication));
        } catch (Exception e) {
            throw new IllegalStateException("Cannot convert access token to JSON", e);
        }
        return JwtHelper.encode(content, signer, customHeaders).getEncoded();
    }

    @Override
    public void setKeyPair(KeyPair keyPair) {
        super.setKeyPair(keyPair);
        this.signer = new RsaSigner((RSAPrivateKey) keyPair.getPrivate());
    }
}
